#include <global_planner/global_planner.h>

// #define GROUND_8

using namespace std;
using namespace ros;
using namespace Eigen;

MavGlobalPlanner::MavGlobalPlanner(Config &conf, NodeHandle &nh_)
    : config(conf), nh(nh_), odomInitialized(false), visualization(config, nh)
{
    triggerSub = nh.subscribe(config.triggerTopic, 1, &MavGlobalPlanner::triggerCallBack, this,
                              TransportHints().tcpNoDelay());
    odomSub = nh.subscribe(config.odomTopic, 1, &MavGlobalPlanner::odomCallBack, this,
                           TransportHints().tcpNoDelay());
    trajPub = nh.advertise<quadrotor_msgs::PolynomialTrajectory>(config.trajectoryTopic, 1);
}

void MavGlobalPlanner::triggerCallBack(const geometry_msgs::PoseStamped::ConstPtr &msg)
{
    if (!odomInitialized)
    {
        return;
    }

    Eigen::Vector3d startTrans = curOdom;
    Eigen::Vector3d xyzMin(config.boxMin[0], config.boxMin[1], config.boxMin[2]);
    Eigen::Vector3d xyzMax(config.boxMax[0], config.boxMax[1], config.boxMax[2]);
    double diagWidth = config.diagWidth;
    double sideWidth = config.sideWidth;
    int round = config.round;

    startTrans(0) = std::max(startTrans(0), xyzMin(0));
    startTrans(1) = std::max(startTrans(1), xyzMin(1));
    startTrans(2) = std::max(startTrans(2), xyzMin(2));
    startTrans(0) = std::min(startTrans(0), xyzMax(0));
    startTrans(1) = std::min(startTrans(1), xyzMax(1));
    startTrans(2) = std::min(startTrans(2), xyzMax(2));

    Eigen::Vector3d goalTrans;

    if(config.trajType == 0)
    {
        // line
        goalTrans = Eigen::Vector3d{config.endWayPoint[0], config.endWayPoint[1], config.endWayPoint[2]};
    }
    else if(config.trajType == 1)
    {
        // circle
        goalTrans = startTrans;
    }
    else if(config.trajType == 2)
    {
        // eight
        goalTrans = (xyzMin + xyzMax) / 2.0;
    }
    
#ifdef GROUND_8
    goalTrans(2) = 0.1;
#endif

    double tanGamma = (startTrans(1) - xyzMin(1)) / (xyzMax(0) - startTrans(0) + FLT_EPSILON);
    double cosGamma = 1.0 / sqrt(1.0 + tanGamma * tanGamma);
    double sinGamma = sqrt(1.0 - cosGamma * cosGamma);
    double xProjIniWidth = 0.5 * diagWidth / sinGamma;
    double yProjIniWidth = 0.5 * diagWidth / cosGamma;

    double tanTheta = (xyzMax(1) - xyzMin(1)) / (xyzMax(0) - xyzMin(0));
    double cosTheta = 1.0 / sqrt(1.0 + tanTheta * tanTheta);
    double sinTheta = sqrt(1.0 - cosTheta * cosTheta);
    double xProjWidth = 0.5 * diagWidth / sinTheta;
    double yProjWidth = 0.5 * diagWidth / cosTheta;

    std::vector<Eigen::MatrixXd> hPolys;
    // hPolys.reserve(50 * round + 1);
   
    
    if(config.trajType == 0)
    {
        // line polyhedra box
        Eigen::MatrixXd poly(6, 6);

        poly.block<3,1>(0,0) =  Eigen::Vector3d(1,0,0);
        poly.block<3,1>(0,1) = -Eigen::Vector3d(1,0,0);
        // 左右面（法向/半径方向）
        poly.block<3,1>(0,2) =  Eigen::Vector3d(0,1,0);
        poly.block<3,1>(0,3) = -Eigen::Vector3d(0,1,0);
        // 上下
        poly.block<3,1>(0,4) = Eigen::Vector3d(0,0,1);
        poly.block<3,1>(0,5) = -Eigen::Vector3d(0,0,1);

        // 面上的点
        poly.block<3,1>(3,0) = goalTrans;
        poly.block<3,1>(3,1) = startTrans;
        poly.block<3,1>(3,2) = diagWidth * Eigen::Vector3d(0,1,0) + Eigen::Vector3d(goalTrans(0)/2, 0, goalTrans(2));
        poly.block<3,1>(3,3) = -diagWidth * Eigen::Vector3d(0,1,0) + Eigen::Vector3d(goalTrans(0)/2, 0, goalTrans(2));
        poly.block<3,1>(3,4) = diagWidth * Eigen::Vector3d(0,0,1) + Eigen::Vector3d(goalTrans(0)/2, 0, goalTrans(2));
        poly.block<3,1>(3,5) = -diagWidth * Eigen::Vector3d(0,0,1) + Eigen::Vector3d(goalTrans(0)/2, 0, goalTrans(2));
        
        hPolys.push_back(poly);
    }
    else if (config.trajType == 1)
    {
        int N = config.boxPieces; // 多面体数量，越大越圆
        double r = config.radius; // 圆半径
        double z = startTrans(2); // 圆所在平面高度
        

        double length_front = diagWidth; // 前后边长度
        double length_back = diagWidth;
        double width_out = 0.25 * diagWidth; // 左右边长度
        double width_in = 0.25 * diagWidth;
        double height_up = 0.125; // 上下边长度
        double height_down = 0.125;
        
        for (int i = 0; i < round; i++) {
        for (int i = 0; i < N; ++i) {
            double theta = 2 * M_PI * i / N;
            Eigen::Vector3d center(startTrans(0) + r - r * cos(-theta), startTrans(1) + r * sin(theta), z);
            Eigen::Vector3d tangent(sin(theta), cos(theta), 0); // 切线方向
            Eigen::Vector3d normal(cos(theta), -sin(theta), 0);   // 半径方向

            Eigen::MatrixXd poly(6, 6);
            // 前后面（切线方向）
            poly.block<3,1>(0,0) =  tangent;
            poly.block<3,1>(0,1) = -tangent;
            // 左右面（法向/半径方向）
            poly.block<3,1>(0,2) =  normal;
            poly.block<3,1>(0,3) = -normal;
            // 上下
            poly.block<3,1>(0,4) = Eigen::Vector3d(0,0,1);
            poly.block<3,1>(0,5) = Eigen::Vector3d(0,0,-1);

            // 面上的点
            poly.block<3,1>(3,0) = center + tangent * length_front;
            poly.block<3,1>(3,1) = center - tangent * length_back;
            poly.block<3,1>(3,2) = center + normal * width_out;
            poly.block<3,1>(3,3) = center - normal * width_in;
            poly.block<3,1>(3,4) = center + Eigen::Vector3d(0,0,height_up);
            poly.block<3,1>(3,5) = center - Eigen::Vector3d(0,0,height_down);

            hPolys.push_back(poly);
        }
        }
    } else if(config.trajType == 2)
    {
         Eigen::MatrixXd box(6, 6), inic(6, 8), lbrt(6, 8), rtlt(6, 7), ltrb(6, 8), rblb(6, 7);

        box << Eigen::Vector3d::UnitX(), Eigen::Vector3d::UnitY(), Eigen::Vector3d::UnitZ(),
            -Eigen::Vector3d::UnitX(), -Eigen::Vector3d::UnitY(), -Eigen::Vector3d::UnitZ(),
                xyzMax, xyzMax, xyzMax, xyzMin, xyzMin, xyzMin;

        inic << box.topRows<3>(), Eigen::Vector3d(sinGamma, cosGamma, 0),
        -Eigen::Vector3d(sinGamma, cosGamma, 0),
        box.bottomRows<3>(), Eigen::Vector3d(xyzMax(0), xyzMin(1) + yProjIniWidth, xyzMax(2)),
        Eigen::Vector3d(xyzMax(0) - xProjIniWidth, xyzMin(1), xyzMax(2));
        
        // polyhedra left-bottom-right-top
        lbrt << box.topRows<3>(), Eigen::Vector3d(sinTheta, cosTheta, 0),
                -Eigen::Vector3d(sinTheta, cosTheta, 0),
        box.bottomRows<3>(), Eigen::Vector3d(xyzMax(0), xyzMin(1) + yProjWidth, xyzMax(2)),
            Eigen::Vector3d(xyzMax(0) - xProjWidth, xyzMin(1), xyzMax(2));

        // polyhedra right-top-left-bottom
        rtlt << box.topRows<3>(), -Eigen::Vector3d::UnitX(),
            box.bottomRows<3>(), xyzMax - sideWidth * Eigen::Vector3d::UnitX();

        // polyhedra left-top-right-bottom
        ltrb << box.topRows<3>(), Eigen::Vector3d(sinTheta, -cosTheta, 0),
            -Eigen::Vector3d(sinTheta, -cosTheta, 0),
            box.bottomRows<3>(), Eigen::Vector3d(xyzMax(0), xyzMax(1) - yProjWidth, xyzMax(2)),
            Eigen::Vector3d(xyzMax(0) - xProjWidth, xyzMax(1), xyzMax(2));

        // polyhedra right-bottom-left-bottom
        rblb << box.topRows<3>(), Eigen::Vector3d::UnitX(),
            box.bottomRows<3>(), xyzMin + sideWidth * Eigen::Vector3d::UnitX();
    
        hPolys.push_back(inic);

        for (int i = 0; i < round; i++)
        {
            hPolys.push_back(rtlt);
            hPolys.push_back(ltrb);
            hPolys.push_back(rblb);
            hPolys.push_back(lbrt);
        }
    }
    
    // configure initial and final states
    Eigen::Matrix3d iniState;
    Eigen::Matrix3d finState;
    iniState << startTrans, Eigen::Vector3d::Zero(), Eigen::Vector3d::Zero();
    finState << goalTrans, Eigen::Vector3d::Zero(), Eigen::Vector3d::Zero();

    double vmax = config.maxVelRate, amax = config.maxAccRate;

    Eigen::Vector3d chi(config.chiVec[0], config.chiVec[1], config.chiVec[2]);

    double res = 2.0;
    int itg = 8;

    std::chrono::high_resolution_clock::time_point tic = std::chrono::high_resolution_clock::now();

    GCOPTER nonlinOpt;
    Trajectory traj;

    if (!nonlinOpt.setup(config.weightT, 1.0, iniState, finState, hPolys, res, itg, vmax, amax, chi, config.c2Diffeo))
    {
        return;
    }

    nonlinOpt.optimize(traj, config.relCostTol);

    std::chrono::high_resolution_clock::time_point toc = std::chrono::high_resolution_clock::now();

    double compTime = std::chrono::duration_cast<std::chrono::microseconds>(toc - tic).count() * 1.0e-3;

    ROS_WARN("\n Time: %.2f s, MaxVel: %.2f m/s, MaxAcc: %.2f m/s^2 \n", traj.getTotalDuration(), traj.getMaxVelRate(), traj.getMaxAccRate());

    if (traj.getPieceNum() > 0)
    {
        quadrotor_msgs::PolynomialTrajectory trajMsg;
        Time stamp = odomStamp;//ros::Time::now();
        polynomialTrajConverter(traj, trajMsg, Eigen::Isometry3d::Identity(), stamp);
        trajPub.publish(trajMsg);
        visualization.visualize(traj, ros::Time::now(), compTime);

        vec_E<Polyhedron3D> polyhedra;
        polyhedra.reserve(config.boxPieces);
        int count = 0;
        for (const auto &ele : hPolys)
        {
            Polyhedron3D hPoly;
            for (int i = 0; i < ele.cols(); i++)
            {
                hPoly.add(Hyperplane3D(ele.col(i).tail<3>(), ele.col(i).head<3>()));
            }
            polyhedra.push_back(hPoly);
            count++;
            if (count > config.boxPieces-1)
            {
                break;
            }
        }
        visualization.visualizePolyH(polyhedra, ros::Time::now());
    }
}

void MavGlobalPlanner::odomCallBack(const nav_msgs::Odometry::ConstPtr &msg)
{
    curOdom(0) = msg->pose.pose.position.x;
    curOdom(1) = msg->pose.pose.position.y;
    curOdom(2) = msg->pose.pose.position.z;
    odomStamp = msg->header.stamp;
    odomInitialized = true;
}

void MavGlobalPlanner::polynomialTrajConverter(const Trajectory &traj,
                                               quadrotor_msgs::PolynomialTrajectory &trajMsg,
                                               Eigen::Isometry3d tfR2L, Time &iniStamp)
{
    trajMsg.header.stamp = iniStamp;
    static uint32_t traj_id = 0;
    traj_id++;
    trajMsg.trajectory_id = traj_id;
    trajMsg.action = quadrotor_msgs::PolynomialTrajectory::ACTION_ADD;
    trajMsg.num_order = traj[0].getOrder();
    trajMsg.num_segment = traj.getPieceNum();
    Eigen::Vector3d initialVel, finalVel;
    initialVel = tfR2L * traj.getVel(0.0);
    finalVel = tfR2L * traj.getVel(traj.getTotalDuration());
    trajMsg.start_yaw = 0.0;
    trajMsg.final_yaw = 0.0;

    for (size_t p = 0; p < (size_t)traj.getPieceNum(); p++)
    {
        trajMsg.time.push_back(traj[p].getDuration());
        trajMsg.order.push_back(traj[p].getCoeffMat().cols() - 1);

        Eigen::VectorXd linearTr(2);
        linearTr << 0.0, trajMsg.time[p];
        std::vector<Eigen::VectorXd> linearTrCoeffs;
        linearTrCoeffs.emplace_back(1);
        linearTrCoeffs[0] << 1;
        for (size_t k = 0; k < trajMsg.order[p]; k++)
        {
            linearTrCoeffs.push_back(RootFinder::polyConv(linearTrCoeffs[k], linearTr));
        }

        Eigen::MatrixXd coefMat(3, traj[p].getCoeffMat().cols());
        for (int i = 0; i < coefMat.cols(); i++)
        {
            coefMat.col(i) = tfR2L.rotation() * traj[p].getCoeffMat().col(coefMat.cols() - i - 1).head<3>();
        }
        coefMat.col(0) = (coefMat.col(0) + tfR2L.translation()).eval();

        for (int i = 0; i < coefMat.cols(); i++)
        {
            double coefx(0.0), coefy(0.0), coefz(0.0);
            for (int j = i; j < coefMat.cols(); j++)
            {
                coefx += coefMat(0, j) * linearTrCoeffs[j](i);
                coefy += coefMat(1, j) * linearTrCoeffs[j](i);
                coefz += coefMat(2, j) * linearTrCoeffs[j](i);
            }
            trajMsg.coef_x.push_back(coefx);
            trajMsg.coef_y.push_back(coefy);
            trajMsg.coef_z.push_back(coefz);
        }
    }

    trajMsg.mag_coeff = 1.0;
    trajMsg.debug_info = "";
}

Visualization::Visualization(Config &conf, NodeHandle &nh_)
    : config(conf), nh(nh_)
{
    routePub = nh.advertise<visualization_msgs::Marker>("/visualization/route", 1);
    wayPointsPub = nh.advertise<visualization_msgs::Marker>("/visualization/waypoints", 1);
    appliedTrajectoryPub = nh.advertise<visualization_msgs::Marker>("/visualization/applied_trajectory", 1);
    hPolyPub = nh.advertise<decomp_ros_msgs::PolyhedronArray>("/visualization/polyhedra", 1);
    textPub = nh.advertise<visualization_msgs::Marker>("/visualization/text", 1);
}

void Visualization::visualize(const Trajectory &appliedTraj, Time timeStamp, double compT)
{
    visualization_msgs::Marker routeMarker, wayPointsMarker, appliedTrajMarker;

    routeMarker.id = 0;
    routeMarker.type = visualization_msgs::Marker::LINE_LIST;
    routeMarker.header.stamp = timeStamp;
    routeMarker.header.frame_id = config.odomFrame;
    routeMarker.pose.orientation.w = 1.00;
    routeMarker.action = visualization_msgs::Marker::ADD;
    routeMarker.ns = "route";
    routeMarker.color.r = 1.00;
    routeMarker.color.g = 0.00;
    routeMarker.color.b = 0.00;
    routeMarker.color.a = 1.00;
    routeMarker.scale.x = 0.05;

    wayPointsMarker = routeMarker;
    wayPointsMarker.type = visualization_msgs::Marker::SPHERE_LIST;
    wayPointsMarker.ns = "waypoints";
    wayPointsMarker.color.r = 1.00;
    wayPointsMarker.color.g = 0.00;
    wayPointsMarker.color.b = 0.00;
    wayPointsMarker.scale.x = 0.20;
    wayPointsMarker.scale.y = 0.20;
    wayPointsMarker.scale.z = 0.20;

    appliedTrajMarker = routeMarker;
    appliedTrajMarker.header.frame_id = config.odomFrame;
    appliedTrajMarker.id = 0;
    appliedTrajMarker.ns = "applied_trajectory";
    appliedTrajMarker.color.r = 0.00;
    appliedTrajMarker.color.g = 0.50;
    appliedTrajMarker.color.b = 1.00;
    appliedTrajMarker.scale.x = 0.10;

    Eigen::MatrixXd route = appliedTraj.getPositions();

    if (route.cols() > 0)
    {
        bool first = true;
        Vector3d last;
        for (int i = 0; i < route.cols(); i++)
        {
            if (first)
            {
                first = false;
                last = route.col(i);
                continue;
            }
            geometry_msgs::Point point;

            point.x = last(0);
            point.y = last(1);
            point.z = last(2);
            routeMarker.points.push_back(point);
            point.x = route.col(i)(0);
            point.y = route.col(i)(1);
            point.z = route.col(i)(2);
            routeMarker.points.push_back(point);
            last = route.col(i);
        }

        routePub.publish(routeMarker);
    }

    if (route.cols() > 0)
    {
        for (int i = 0; i < route.cols(); i++)
        {
            geometry_msgs::Point point;
            point.x = route.col(i)(0);
            point.y = route.col(i)(1);
            point.z = route.col(i)(2);
            wayPointsMarker.points.push_back(point);
        }

        wayPointsPub.publish(wayPointsMarker);
    }

    if (appliedTraj.getPieceNum() > 0)
    {
        double T = 0.01;
        Vector3d lastX = appliedTraj.getPos(0.0);
        for (double t = T; t < appliedTraj.getTotalDuration(); t += T)
        {
            geometry_msgs::Point point;
            Vector3d X = appliedTraj.getPos(t);
            point.x = lastX(0);
            point.y = lastX(1);
            point.z = lastX(2);
            appliedTrajMarker.points.push_back(point);
            point.x = X(0);
            point.y = X(1);
            point.z = X(2);
            appliedTrajMarker.points.push_back(point);
            lastX = X;
        }
        appliedTrajectoryPub.publish(appliedTrajMarker);
    }

    if (appliedTraj.getPieceNum() > 0)
    {

        visualization_msgs::Marker textMarker;
        textMarker.header.frame_id = config.odomFrame;
        textMarker.header.stamp = timeStamp;
        textMarker.ns = "text";
        textMarker.id = 1;
        textMarker.type = visualization_msgs::Marker::TEXT_VIEW_FACING;
        textMarker.action = visualization_msgs::Marker::ADD;

        textMarker.pose.position.x = 16.5;
        textMarker.pose.position.y = -11.0;
        textMarker.pose.position.z = 6.0;
        textMarker.pose.orientation.x = 0.0;
        textMarker.pose.orientation.y = 0.0;
        textMarker.pose.orientation.z = 0.0;
        textMarker.pose.orientation.w = 1.0;
        textMarker.scale.x = 1.0;
        textMarker.scale.y = 1.0;
        textMarker.scale.z = 1.0;
        textMarker.color.r = 1.0;
        textMarker.color.g = 0.0;
        textMarker.color.b = 0.0;
        textMarker.color.a = 1.0;
        textMarker.text = "Comp:";
        textMarker.text += to_string((int)(compT));
        textMarker.text += ".";
        textMarker.text += to_string((int)(compT * 10) % 10);
        textMarker.text += "ms";

        textPub.publish(textMarker);
    }
}

void Visualization::visualizePolyH(const vec_E<Polyhedron3D> &polyhedra, ros::Time timeStamp)
{
    decomp_ros_msgs::PolyhedronArray poly_msg = DecompROS::polyhedron_array_to_ros(polyhedra);
    poly_msg.header.frame_id = config.odomFrame;
    poly_msg.header.stamp = timeStamp;
    hPolyPub.publish(poly_msg);
}